import numpy as np
import torch
import os
import torch.backends.cudnn as cudnn
from torchvision import transforms
import torchvision
from PIL import Image
from tqdm import tqdm
from model.PR_Grim_EffiNet import PR_EffiNet
from model.vision_transformer import ViT
from model.efficientNetv2 import effnetv2_s
from model.efficientNetv1 import EfficientNet
from model.ghostNet import ghostnet
import config

test_image_path = config.test_data_path            #图片文件夹路径，或图片路径
input_shape = config.image_shape
checkpoint_path = config.test_checkpoint_path
select_model = config.selected_model
cuda = config.cuda
# test_image_path = "/media/liyan/3b75ab20-92d3-4be6-81c2-4e1798e2fe16/我的数据集/图像分类数据集/new1/test/"
# checkpoint_path = "/media/liyan/3b75ab20-92d3-4be6-81c2-4e1798e2fe16/权重/GEEiNetV2/B2_checkpoint/48/checkpoint_57.pth"
# input_shape = [3,48,48]
# cuda = False
# select_model = "B2"             #模型大小选择,['B0','B1','B2','B3','B4','B5','effiv2','effiv1','ghostnet']

num_workers = 0
drop_last = True
batch_size = 56
caculate_confused_matrix = False #是否需要计算混淆矩阵
#-----------------------------一些函数--------------------------------------
def get_dataset_length(path:str)->int:
    count = 0
    for dir_class in os.listdir(path):
        count += len(os.listdir(os.path.join(path,dir_class)))
    return count

def sort_file(file_name:str):
    if file_name.endswith('.txt'):
        return float(0)
    else:
        ret = file_name.split(".")[0].split("_")[1]
        return float(ret)

def get_class_names()->list:
    return os.listdir(test_image_path)

def test_datasets(filename:str,index:int,model,data_sum:int,test_data_loader):
    model.eval()
    top1_acc_sum = 0
    top5_acc_sum = 0
    confued_matrix = [[],[]]    #前者是x，预测值，后者是y，真实值
    with tqdm(desc=f"index:{index},filename:{filename}",postfix=dict,ncols=250) as pbar:
        for step, (inputs, labels) in enumerate(test_data_loader, 0):
            with torch.no_grad():#被这句话包住的代码，不会跟踪反向梯度计算，不会生成计算图，不会计算梯度，更不会进行反向传播，
                if cuda:
                    inputs = torch.autograd.Variable(inputs).cuda()
                    labels = torch.autograd.Variable(labels).cuda()
                else:
                    inputs = torch.autograd.Variable(inputs)
                    labels = torch.autograd.Variable(labels)

                outputs = model(inputs)
                #---------Top1准确率--------------
                _, val_pre = torch.max(outputs, 1)
                top1_acc_sum += torch.eq(val_pre, labels).sum()
                confued_matrix[0].extend(val_pre.cpu().numpy())
                confued_matrix[1].extend(labels.cpu().numpy())


                #---------Top5准确率---------------
                max5 = max((1,5))
                y_resize = labels.view(-1,1)
                _,pred = outputs.topk(max5,1,largest=True,sorted=True)
                top5_acc_sum += torch.eq(pred,y_resize).sum().float().item()

                pbar.set_postfix(**{'Top1':float(top1_acc_sum) / float(data_sum),'Top5':float(top5_acc_sum)/float(data_sum)})
                pbar.update(1)

    return float(top1_acc_sum) / float(data_sum),float(top5_acc_sum)/float(data_sum),np.asarray(confued_matrix)

if __name__ == "__main__":
    class_names = get_class_names()
    #----------------------读取数据-----------------------------------.
    transform = torchvision.transforms.Compose([
        # transforms.CenterCrop(10),
        transforms.Resize((input_shape[1], input_shape[2])),  # 重置图像分辨率,会降低分辨率
        # transforms.RandomHorizontalFlip(),  # 根据某种概率水平翻转(default:p=0.5),关于左右对称轴左右对称
        # transforms.Grayscale(input_shape[0]),  # 转换为灰度图,3就是灰色，1就是加绿色(default)
        # transforms.RandomRotation(2),  # 随机旋转，表示随机旋转+-角度范围
        # transforms.ColorJitter(brightness=,contrast=,saturation=,hue=)#亮度(>0)，对比度(>0)，饱和度(>0)，色相(-0.5,0.5)
        # transforms.LinearTransformation(),#线性变换，可用于白化处理
        # transforms.RandomCrop(),    #随机选择区域裁剪
        # transforms.RandomResizedCrop(),#随机大小，随机长宽比裁剪原始图片，最后将图片resize到设定好的size
        # transforms.FiveCrop(),      #上下左右中心裁剪，获得5张图片，返回一个4D-tensor参数
        # transforms.TenCrop(),       #上下左右中心裁剪翻转，获得10张图片
        # transforms.RandomVerticalFlip(),# 根据某种概率垂直翻转(default:p=0.5),关于垂直对称轴左右对称
        # transforms.Normalize(),     #按通道进行标准化，先减均值再除标准差
        # transforms.Pad(),#填充
        # transforms.RandomAffine(),  #仿射变换,可理解投影
        # transforms.RandomGrayscale(),#依据概率p变成灰度图
        # transforms.RandomChoice(),#随机在给定的一系列transform中选择一个进行操作、
        # transforms.RandomApply(),#给一个transform加概率，以某种概率进行操作
        # transforms.RandomOrder(),#对transforms的操作随机打乱
        transforms.ToTensor()  # 将PIL Image或numpy.ndarry转换成torch.Tensor
    ])
    test_data = torchvision.datasets.ImageFolder(root=test_image_path,transform=transform)
    test_data_loader = torch.utils.data.DataLoader(dataset=test_data,batch_size=batch_size,shuffle=True,
                                                   num_workers=num_workers,drop_last=drop_last)
    count_test = get_dataset_length(path=test_image_path) - get_dataset_length(path=test_image_path) % batch_size
    print(f"测试集数量：{count_test}")
    # ----------------------------模型读入----------------------------------------------
    num_class = 43
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    if select_model == 'effiv2':
        model = effnetv2_s(num_classes=num_class)
        print('selected model is :', select_model)
    elif select_model == 'effiv1':
        model = EfficientNet(width_coeff=1, depth_coeff=1, num_classes=num_class)
        print('selected model is :', select_model)
    elif select_model == 'ghostnet':
        model = ghostnet(num_classes=num_class)
        print('selected model is :', select_model)
    elif select_model == 'vit':
        model = ViT(input_shape[1],num_classes=num_class,patch_size=input_shape[1]//16)
        print('selected model is :', select_model)
    else:
        model = PR_EffiNet(select_model=select_model, num_classes=num_class, activation='relu')
        print('selected model is :', select_model)

    #------------------设置评判准则-----------------------------
    crierion = torch.nn.CrossEntropyLoss()
    #--------------------------放入cuda-----------------------
    if cuda:
        if torch.cuda.device_count() > 1:
            model = torch.nn.DataParallel(model)
        cudnn.benchmark = True
        model = model.cuda()
    #----------------------加载模型----------------------------
    device = torch.device("cuda:0" if cuda else "cpu")

    if cuda:
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.enabled = True
        torch.backends.cudnn.benchmark = True
    #-----------------对每个权重做测试-----------------
    if os.path.isdir(checkpoint_path) and caculate_confused_matrix == False:
        checkpoint_file = os.listdir(checkpoint_path)
        for index,fila_name in enumerate(checkpoint_file):
            txt_file = open(checkpoint_path+"test_result.txt",'a+')

            if fila_name.endswith('.txt'):
                continue

            checkpoint_file_path = os.path.join(checkpoint_path,fila_name)
            checkpoint = torch.load(checkpoint_file_path,map_location=device)
            model.load_state_dict(checkpoint['net'])
            top1,top5,confused_matrix = test_datasets(fila_name,index,model,data_sum=count_test,test_data_loader=test_data_loader)
            txt_file.write("Filename:%-20s|Top1:%-10.5f|Top5:%-10.5f\n"%(fila_name,top1,top5))
            txt_file.close()
    #----------------对单个权重做测试-------------
    elif os.path.isfile(checkpoint_path):
        txt_file = open(os.path.abspath(checkpoint_path)+"the_best_val_pth.txt",'a+')
        checkpoint_dict = torch.load(checkpoint_path,map_location=device)
        model.load_state_dict(checkpoint_dict['net'])
        top1,top5,confued_matrix = test_datasets(os.path.basename(checkpoint_path),0,model,count_test,test_data_loader)
        txt_file.write("Filename:%-20s|Top1:%-10.5f|Top5:%-10.5f\n"%(os.path.basename(checkpoint_path),top1,top5))
        txt_file.close()

        #--------计算混淆矩阵--------------------
        if caculate_confused_matrix:
            from confused_matrix import plot_confusion_matrix
            plot_confusion_matrix(confused_matrix=confued_matrix,num_classes=num_class)


